/**
 * @file layerEarth_MT.h
 * @author Yi Zhang (zhangyiss@icloud.com)
 * @brief A closed-form for plane wave electromagnetic fields (PWEF) on a n-layers Earth model with an oblique angle at a given frequency
 * @version 1.0
 * @date 2021-11-30
 * 
 * @copyright Copyright (c) 2021
 * 
 * +z downward, z=0 the Air-Earth interface
 * time dependence exp(-iwt)
 * TE-mode: E = (Ex,0 ,0 ), H = (0, Hy,Hz)
 * TM-mode: E = (0 ,Ey,Ez), H = (Hx,0 ,0 )
 *          x (Ex for TE mode and Hx for TM mode)
 *        /
 *      /
 *    /    air (layer 0)
 *  /
 * O---------------------> y
 * |   layer 1
 * |---------------------
 * |   layer 2
 * |---------------------
 * |   ...
 * |---------------------
 * |   layer n
 * |
 * z
 * 
 */

#ifndef _LAYEREARTH_MT_H
#define _LAYEREARTH_MT_H

#include "vector"
#include "complex"

struct matrix2x2
{
    std::complex<double> val[2][2];
};

struct conductive_layer
{
    double cndt; // conductivity (S/m)
    double epsl; // epsilon (relative dielectric constant)
    double mu; // relative permeability
    double btm_depth; // bottom depths of the layer. this is also the top depth for the next layer

    conductive_layer(){}
    conductive_layer(double c, double e, double m, double d);
    void set(double c, double e, double m, double d);
};

class layerEarthMT
{
public:
    layerEarthMT();
    layerEarthMT(const std::vector<conductive_layer> &layers);
    virtual ~layerEarthMT(){}

    void initialize();
    void calculate_EH(double z, double y);

    std::complex<double> get_Ex();
    std::complex<double> get_Ey();
    std::complex<double> get_Ez();
    std::complex<double> get_Hx();
    std::complex<double> get_Hy();
    std::complex<double> get_Hz();
    std::complex<double> get_Ex_dy();
    std::complex<double> get_Ex_dz();
    std::complex<double> get_Hx_dy();
    std::complex<double> get_Hx_dz();
    double get_apparent_resistivity();
    double get_phase_angle();

    void set_layers(const std::vector<conductive_layer> &layers);
    void set_frequency(double frequency);
    void set_incident_angle(double theta);
    void set_to_TM_mode();
    void set_to_TE_mode();
    void set_ground_amplitude(double b0);

protected:
    bool initialized_;
    double omega_; // angular frequency (default is 2*pi, aka frequency = 1.0 Hz)
    double theta_; // incident angle (default is 0)
    bool TM_mode_; // if this is true then works in TM mode, otherwise works in TE mode (default is false)
    double B0_; // amplitude at ground (default is 1.0)
    std::vector<conductive_layer> all_layers_;

    std::complex<double> ky_;
    std::complex<double> Ex_, Ey_, Ez_, Hx_, Hy_, Hz_;
    std::complex<double> Ex_y_, Ex_z_, Hx_y_, Hx_z_;
    std::vector<double> h_; // layer thicknesses
    std::vector<std::complex<double> > k_, kz_, zhat_, yhat_; // wave number of each layer
    std::vector<std::complex<double> > Z_, Y_, U_, D_;
    std::vector<matrix2x2> T_; // transform matrices
};

#endif // _LAYEREARTH_MT_H